# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     train
   Description :   
   Author :       lth
   date：          2022/12/9
-------------------------------------------------
   Change Activity:
                   2022/12/9 9:14: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import cv2
import numpy as np
import torch
from torch import optim
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import GetConfig
from datalist import rescale, CelebAHQ
from model import NICEModel
from utils import GaussianPriorNICELoss

total_loss = 999999999999999999999999999


class Train:
    def __init__(self):
        self.args = GetConfig()
        print(f"-----------{self.args.project_name}-------------")
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if use_cuda else "cpu")

        kwargs = {"num_workers": 2, "pin_memory": True} if use_cuda else {"num_workers": 2, "pin_memory": False}

        # self.train_dataset, self.zca = load_cifar10(train=True)
        # self.test_dataset, _ = load_cifar10(train=False)

        self.train_dataset = CelebAHQ(mode="train")
        self.test_dataset = CelebAHQ(mode="test")
        self.train_dataloader = DataLoader(self.train_dataset, batch_size=80, shuffle=True,
                                           **kwargs)
        self.test_dataloader = DataLoader(self.test_dataset, batch_size=2, shuffle=True,
                                          **kwargs)

        self.model = NICEModel(64 * 64 * 3, 1000, 5).to(self.device)

        if self.args.resume:
            try:
                print("\nload the weight from pretrained-weight file")
                model_dict = self.model.state_dict()
                checkpoint = torch.load("weights/best.pth")['model_state_dict']
                model_dict.update(checkpoint)
                self.model.load_state_dict(model_dict, strict=True)
                print("Restoring the weight from pretrained-weight file \nFinished loading the weight\n")
            except Exception as e:
                raise e

        if use_cuda:
            cudnn.benchmark = True
            cudnn.enabled = True

        self.criterion = GaussianPriorNICELoss()

        self.optimizer = optim.Adam(
            lr=self.args.lr,
            params=[
                {"params": self.model.parameters()}],
            betas=(0.9, 0.999),
            weight_decay=1e-4
        )
        self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.99)

    def work(self):
        for epoch in range(1, self.args.epochs):
            self.train(epoch)
            if epoch % 10 == 0:
                self.inference()
            self.scheduler.step()
        torch.cuda.empty_cache()
        print("model finish training")

    def train(self, epoch):
        self.model.train()
        average_loss = []

        pbar = tqdm(self.train_dataloader, desc=f'Train Epoch:{epoch}/{self.args.epochs}')

        for image in pbar:
            image = image.view(image.shape[0], -1)
            image = image + torch.rand_like(image).div_(128.).add_(-1. / 256.)
            image = rescale(image, -1, 1)

            self.optimizer.zero_grad()
            image = image.to(self.device)
            output = self.model(image)

            loss = self.criterion(output, self.model.scaling_diag)
            loss.backward()
            average_loss.append(loss.item())
            self.optimizer.step()

            pbar.set_description(
                f'Train Epoch: {epoch}/{self.args.epochs} '
                f' train_loss: {np.mean(average_loss)} '
                f' learning_rate: {self.optimizer.state_dict()["param_groups"][0]["lr"]}'

            )
        global total_loss
        if np.mean(average_loss) < total_loss:
            torch.save({
                'model_state_dict': self.model.state_dict(),
            },
                './weights/best.pth')
            print("model saved")
            total_loss = np.mean(average_loss)

    @torch.no_grad()
    def inference(self):

        self.model.eval()
        imgs_row = []
        for i in range(10):
            imgs_col = []
            for j in range(10):

                y = torch.randn(1, 64 * 64 * 3).to(self.device)
                x = self.model.inverse(y)

                x = x.view(1, 3, 64, 64)
                epslion = 1e-5
                x = torch.clip(x, -1 + epslion, 1 - epslion)
                x = ((x + 1) / 2 * 255).cpu().numpy().astype(np.uint8)[0]
                x = np.transpose(x, [1, 2, 0])[..., ::-1]
                imgs_col.append(x.copy())
            imgs_row.append(np.concatenate(imgs_col,axis=0))
        img = np.concatenate(imgs_row,axis=1)

        cv2.imwrite("result.jpg", img)


if __name__ == "__main__":
    model = Train()
    model.work()
